import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from PIL import Image, ImageTk
import os
import shutil

class YOLOLabelTool:
    def __init__(self, master):
        self.master = master
        self.master.title("YOLO 标注工具")
        self.master.geometry("1400x900")
        self.master.bind("<Control-s>", self.save_labels_shortcut)

        self.image_paths = []
        self.current_index = 0
        self.image = None
        self.tk_image = None
        self.bboxes = []
        self.all_classes = []    # 所有类别，从 classes.txt 读入
        self.filtered_classes = []  # 过滤后的类别
        self.class_var = tk.StringVar()

        self.hero_classes = []
        self.selected_heroes = []

        self.start_x = self.start_y = 0
        self.rect_id = None

        # 左侧操作栏，默认显示
        self.left_frame = tk.Frame(master, width=220, bg="lightgray")
        self.left_frame.pack(side=tk.LEFT, fill=tk.Y)

        # 左侧内容
        tk.Label(self.left_frame, text="操作菜单", bg="lightgray", font=("Arial", 14)).pack(pady=10)

        tk.Button(self.left_frame, text="加载文件夹", width=18, command=self.load_folder).pack(pady=4)

        self.selected_hero_frame = tk.Frame(self.left_frame, bg="lightgray")
        self.selected_hero_frame.pack(pady=(15, 5), fill=tk.X)

        tk.Label(self.left_frame, text="英雄名称", bg="lightgray").pack(pady=(2, 0))
        self.hero_combobox = ttk.Combobox(self.left_frame, height=15)
        self.hero_combobox.pack(pady=2)
        self.hero_combobox.bind("<<ComboboxSelected>>", self.select_hero)
        self.load_heroes()

        # 隐藏菜单按钮，放最下面
        btn_hide_menu = tk.Button(self.left_frame, text="隐藏菜单", width=18, command=self.hide_left_menu)
        btn_hide_menu.pack(side=tk.BOTTOM, pady=10)

        # 右侧主区域，包含图片画布和下方操作栏
        right_frame = tk.Frame(master)
        right_frame.pack(side=tk.RIGHT, expand=True, fill=tk.BOTH)

        self.canvas = tk.Canvas(right_frame, cursor="tcross", bg="white")
        self.canvas.pack(expand=True, fill=tk.BOTH)

        self.canvas.bind("<ButtonPress-1>", self.on_mouse_down)
        self.canvas.bind("<B1-Motion>", self.on_mouse_drag)
        self.canvas.bind("<ButtonRelease-1>", self.on_mouse_up)

        control_frame = tk.Frame(right_frame)
        control_frame.pack(fill=tk.X, pady=5)

        # 打开菜单按钮
        btn_toggle_menu = tk.Button(control_frame, text="打开菜单", width=10, command=self.toggle_left_menu)
        btn_toggle_menu.pack(side=tk.LEFT, padx=5)

        # 上一张 / 下一张
        btn_prev = tk.Button(control_frame, text="上一张", width=10, command=self.prev_image)
        btn_prev.pack(side=tk.LEFT, padx=5)
        btn_next = tk.Button(control_frame, text="下一张", width=10, command=self.next_image)
        btn_next.pack(side=tk.LEFT, padx=5)

        # 类别选择下拉框（只显示过滤后类别）
        tk.Label(control_frame, text="英雄类别:").pack(side=tk.LEFT, padx=(20, 5))
        self.class_dropdown = ttk.Combobox(control_frame, textvariable=self.class_var, width=30)
        self.class_dropdown.pack(side=tk.LEFT)
        # 类别过滤不需要单独响应了，选择英雄后刷新类别

        # 保存按钮
        btn_save_train = tk.Button(control_frame, text="保存标注 (Ctrl+S)", command=self.save_labels)
        btn_save_train.pack(side=tk.RIGHT, padx=10)
        btn_save_val = tk.Button(control_frame, text="保存为验证标注", command=self.save_labels_val)
        btn_save_val.pack(side=tk.RIGHT)

        self.status_label = tk.Label(master, text="", fg="green", font=("Arial", 11))
        self.status_label.pack(side=tk.BOTTOM, pady=5)

        # 读取所有类别
        self.load_all_classes()

    def hide_left_menu(self):
        self.left_frame.pack_forget()

    def toggle_left_menu(self):
        if self.left_frame.winfo_ismapped():
            self.left_frame.pack_forget()
        else:
            self.left_frame.pack(side=tk.LEFT, fill=tk.Y)

    def load_heroes(self):
        if os.path.exists("labels/heros.txt"):
            with open("labels/heros.txt", "r", encoding="utf-8") as f:
                self.hero_classes = [line.strip() for line in f.readlines()]
                self.hero_combobox["values"] = self.hero_classes

    def load_all_classes(self):
        if os.path.exists("labels/classes.txt"):
            with open("labels/classes.txt", "r", encoding="utf-8") as f:
                self.all_classes = [line.strip() for line in f.readlines()]

    def select_hero(self, event=None):
        hero = self.hero_combobox.get().strip()
        if not hero:
            return
        if hero not in self.selected_heroes:
            self.selected_heroes.append(hero)
            self.update_selected_hero_labels()
        # 选择英雄后，刷新类别列表，只显示前缀匹配的类别
        self.refresh_class_list()

    def update_selected_hero_labels(self):
        for widget in self.selected_hero_frame.winfo_children():
            widget.destroy()

        row = 0
        for i, hero in enumerate(self.selected_heroes):
            col = i % 2
            if col == 0:
                row += 1
            frame = tk.Frame(self.selected_hero_frame, bg="lightgray")
            frame.grid(row=row, column=col, padx=2, pady=2)
            label = tk.Label(frame, text=hero, bg="#dfe", relief="solid")
            label.pack(side="left")
            btn = tk.Button(frame, text="❌", width=2, command=lambda h=hero: self.remove_hero(h))
            btn.pack(side="right")

    def remove_hero(self, hero):
        if hero in self.selected_heroes:
            self.selected_heroes.remove(hero)
            self.update_selected_hero_labels()
            self.refresh_class_list()

    def load_folder(self):
        folder = filedialog.askdirectory()
        if not folder:
            return
        self.image_paths = sorted([os.path.join(folder, f) for f in os.listdir(folder)
                                   if f.lower().endswith(('.jpg', '.png', '.jpeg'))])
        self.current_index = 0
        self.load_image()

    def load_image(self):
        if not self.image_paths:
            return

        self.bboxes = []
        path = self.image_paths[self.current_index]
        original_image = Image.open(path)

        # 获取 canvas 的大小，确保 update 后拿到真实尺寸
        self.master.update_idletasks()
        canvas_width = self.canvas.winfo_width()
        canvas_height = self.canvas.winfo_height()

        # 默认大小避免初次打开时 canvas 过小
        if canvas_width <= 1 or canvas_height <= 1:
            canvas_width, canvas_height = 800, 600

        # 保存原图尺寸，用于标注时计算比例
        self.original_image_size = original_image.size

        # 根据 Pillow 版本处理 Resampling 模式
        try:
            resample = Image.Resampling.LANCZOS
        except AttributeError:
            resample = Image.ANTIALIAS

        # 创建缩放后的图片副本，适应 canvas 尺寸
        self.image = original_image.copy()
        self.image.thumbnail((canvas_width, canvas_height), resample)

        self.tk_image = ImageTk.PhotoImage(self.image)
        self.canvas.delete("all")

        # 计算使图片居中的位置
        x = (canvas_width - self.image.width) // 2
        y = (canvas_height - self.image.height) // 2
        self.image_offset = (x, y)  # 用于坐标转换

        self.canvas.create_image(x, y, anchor="nw", image=self.tk_image)
        self.master.title(f"YOLO 标注工具 - {os.path.basename(path)}")

    def prev_image(self):
        if self.current_index > 0:
            self.current_index -= 1
            self.load_image()

    def next_image(self):
        if self.current_index < len(self.image_paths) - 1:
            self.current_index += 1
            self.load_image()
        else:
            messagebox.showinfo("提示", "已经是最后一张图片了")

    def on_mouse_down(self, event):
        self.start_x, self.start_y = event.x, event.y
        self.rect_id = self.canvas.create_rectangle(self.start_x, self.start_y, self.start_x, self.start_y, outline='red')

    def on_mouse_drag(self, event):
        if self.rect_id:
            self.canvas.coords(self.rect_id, self.start_x, self.start_y, event.x, event.y)

    def on_mouse_up(self, event):
        if not self.class_var.get():
            messagebox.showerror("错误", "请先选择一个英雄类别")
            if self.rect_id:
                self.canvas.delete(self.rect_id)
            return
        x0, y0 = self.start_x, self.start_y
        x1, y1 = event.x, event.y
        x0, x1 = sorted([x0, x1])
        y0, y1 = sorted([y0, y1])
        cls = self.class_var.get()
        self.bboxes.append((x0, y0, x1, y1, cls))
        if self.rect_id:
            self.canvas.delete(self.rect_id)
        self.canvas.create_rectangle(x0, y0, x1, y1, outline='green')
        self.canvas.create_text(x0 + 5, y0 + 10, text=cls, fill='green', anchor="nw")
        self.rect_id = None

    def save_labels_shortcut(self, event=None):
        self.save_labels()

    def save_labels(self, prefix_dir= "/Users/shrek/Documents/proj/train"):
        if not self.image_paths:
            return
        label_dir = os.path.join(prefix_dir, "labels")
        image_dir = os.path.join(prefix_dir, "images")
        os.makedirs(label_dir, exist_ok=True)
        os.makedirs(image_dir, exist_ok=True)

        img_w, img_h = self.image.size
        base = os.path.splitext(os.path.basename(self.image_paths[self.current_index]))[0]

        # 保存标注txt文件
        with open(os.path.join(label_dir, base + ".txt"), "w", encoding="utf-8") as f:
            for x0, y0, x1, y1, cls in self.bboxes:
                # 先获取类别索引
                if cls in self.all_classes:
                    cls_idx = self.all_classes.index(cls)
                else:
                    cls_idx = 0
                # 转成相对坐标 (YOLO格式)
                cx = (x0 + x1) / 2 / img_w
                cy = (y0 + y1) / 2 / img_h
                w = (x1 - x0) / img_w
                h = (y1 - y0) / img_h
                f.write(f"{cls_idx} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")

        # 复制图片
        shutil.copy(self.image_paths[self.current_index], os.path.join(image_dir, os.path.basename(self.image_paths[self.current_index])))

        self.status_label.config(text=f"保存成功：{base}.txt")
        #messagebox.showinfo("保存成功", f"标签已保存至{label_dir}")

    def save_labels_val(self):
        self.save_labels(prefix_dir="/Users/shrek/Documents/proj/val")

    def refresh_class_list(self):
        # 根据已选英雄过滤类别
        if not self.selected_heroes:
            # 如果没有选择英雄，显示空类别列表
            self.filtered_classes = []
        else:
            filtered = []
            for hero in self.selected_heroes:
                filtered.extend([cls for cls in self.all_classes if cls.startswith(hero)])
            # 去重
            self.filtered_classes = list(dict.fromkeys(filtered))

        self.class_dropdown['values'] = self.filtered_classes
        if self.filtered_classes:
            self.class_var.set(self.filtered_classes[0])
        else:
            self.class_var.set('')

def main():
    root = tk.Tk()
    app = YOLOLabelTool(root)
    root.mainloop()

if __name__ == '__main__':
    main()